#include "arch_context.h"
#include <arch/arch.h>
#include <task/task.h>

extern uint64_t clone_page_table(uint64_t src_root_phys, uint64_t stack_start, uint64_t stack_end);

/* arch_context_init: 初始化 arch_context（只保留 pt_regs 与 page_table） */
void arch_context_init(arch_context_t *context, uint64_t page_table_addr, uint64_t entry, uint64_t stack, bool user_mode)
{
    memset(context, 0, sizeof(arch_context_t));

    /* 为 satp/页表根物理地址分配一个 slot（保持 header 中 uint64_t *satp 的类型） */
    context->satp = (uint64_t *)malloc(sizeof(uint64_t));
    if (!context->satp) {
        printk("arch_context_init: malloc satp slot failed\n");
        return;
    }
    *(context->satp) = page_table_addr; /* 存放物理页表根（或根 PA），clone_page_table 使用相同语义 */

    /* 在给定栈上为 pt_regs 留出空间（约定 ret_from_trap_handler 从 sp 恢复） */
    context->pt_regs = (struct pt_regs *)stack - 1;
    memset(context->pt_regs, 0, sizeof(struct pt_regs));

    /* 设置最小必须字段：sepc (epc)、sp、a0 清零 */
    context->pt_regs->epc = entry;
    context->pt_regs->sp = stack;
    context->pt_regs->a0 = 0;

    /* sstatus: 设置 UXL=2(64bit) 和 SPIE=1；若为内核线程则同时打开 SIE（和你之前实现一致） */
    if (user_mode) {
        context->pt_regs->sstatus = (2UL << 32) | (1UL << 5); /* UXL=2, SPIE=1 */
    } else {
        context->pt_regs->sstatus = (2UL << 32) | (1UL << 8) | (1UL << 5) | (1UL << 1); /* + SIE */
    }
}

/* arch_context_copy: 为 dst 创建新页表，复制寄存器快照 */
void arch_context_copy(arch_context_t *dst, arch_context_t *src, uint64_t stack)
{
    if (!dst || !src) return;


    
    dst->satp = (uint64_t *)malloc(sizeof(uint64_t));
    

    if (!dst->satp) {
        printk("arch_context_copy: malloc satp failed\n");
        return;
    }

    if (!src->satp) {
        printk("arch_context_copy: src->satp == NULL\n");
        free(dst->satp);
        dst->satp = NULL;
        return;
    }
    uint64_t satp_val = (src->satp);
    /* 打印原始 satp */
    printk("[arch_context_copy] src->satp=0x%lx  \n", satp_val);

    /* 粗暴提取物理地址：去掉高 16 位（MODE+ASID 部分） */
    uint64_t root_phys = (satp_val & ((1ULL << 44) - 1)) << 12;

    /* 打印转换结果 */
    printk("[arch_context_copy] extracted root_phys=0x%lx\n", root_phys);
    
    uint64_t new_root = clone_page_table(root_phys, USER_STACK_START, USER_STACK_END);
    if (!new_root) {
        printk("arch_context_copy: clone_page_table failed\n");
        free(dst->satp);
        dst->satp = NULL;
        return;
    }
    *(dst->satp) = new_root;

    /* 复制 pt_regs：dst->pt_regs 必须已在调用方为目标线程的栈分配好空间 */
    /* 如果 dst->pt_regs 尚未指向合法的栈位置，调用者需先分配（通常在 create_thread 时已经做了） */
    if (dst->pt_regs)
        memcpy(dst->pt_regs, src->pt_regs, sizeof(struct pt_regs));
    else
        printk("arch_context_copy: dst->pt_regs == NULL (确保已为目标线程分配内核栈并设置 pt_regs 指针)\n");
}

void arch_context_free(arch_context_t *context)
{   
    //riscv fpu is failed
    // if (context->fpu_ctx)
    // {
    //     free_frames(virt_to_phys((uint64_t)context->fpu_ctx), 1);
    // }
}

/* 使用 tp 保存 current task 指针（RISC-V 约定） */
task_t *arch_get_current()
{
    task_t *current = NULL;
    asm volatile("mv %0, tp\n\t" : "=r"(current));
    return current;
}

void arch_set_current(task_t *current)
{
    asm volatile("mv tp, %0\n\t" ::"r"(current));
    // printk("arch_set_current -> tp=%p\n", current);
}

/* 切换页表 + 设置 sscratch (kernel_stack) + 把 next->pt_regs 放到 sp 然后跳转到 ret_from_trap_handler 恢复 */
void arch_switch_with_context(arch_context_t *prev, arch_context_t *next, uint64_t kernel_stack)
{
    /* 写 sscratch 供 trap 恢复/中断用（若你的 trap handler 使用 sscratch）*/
    csr_write(sscratch, kernel_stack);

    /* 组装并写入 satp：请确认 MAKE_SATP_PADDR 的语义（是否需要 >>12） */
    uint64_t satp_val = MAKE_SATP_PADDR(SATP_MODE_SV48, 0, *(next->satp));
    write_satp(satp_val);

    /* 全局 TLB flush */
    asm volatile("sfence.vma zero, zero\n\t" ::: "memory");

    /* 切换内核栈到 next 的 pt_regs（ret_from_trap_handler 会根据 pt_regs 恢复并 sret） */
    asm volatile(
        "mv sp, %0\n\t"
        "j ret_from_trap_handler\n\t"
        :
        : "r"(next->pt_regs)
        : "memory");
}

/* 上层调用：pt_regs 是 trap handler 已经保存到当前内核栈的 pt_regs 指针 */
void arch_task_switch_to(struct pt_regs *pt_regs, task_t *prev, task_t *next)
{   
    
    /* 保存当前 trap 的 regs 到 prev 的 arch_context（以便下次调度回来时恢复） */
    if (prev && prev->arch_context)
        prev->arch_context->pt_regs = pt_regs;

    /* 如果同一任务则无需切换 */
    if (prev == next) return;

    /* 简单的防护：如果 next 为 NULL，直接返回（上层应保证存在就绪任务） */
    if (!next || !next->arch_context) return;


    /* ✅ 确保下一个任务的 pt_regs 里 tp 已正确设置 */
    if (next->arch_context->pt_regs) {
        next->arch_context->pt_regs->tp = (uint64_t)next;
    }

    

    /* 状态变更（上层 task 结构的名字可能不同，保持你的字段） */
    if (prev) prev->state = TASK_READY;
    next->state = TASK_RUNNING;

    /* 切换 current pointer */
    arch_set_current(next);
    /* 执行低级切换（页表 + 栈 + 恢复） */
    arch_switch_with_context(prev ? prev->arch_context : NULL, next->arch_context, next->kernel_stack);

    /* 当/如果执行流再次返回到这里，说明本核再次恢复到该点（或别处） */
}

void arch_context_to_user_mode(arch_context_t *context, uint64_t entry, uint64_t stack)
{
    
    if (!context || !context->pt_regs) 
    {
        printk("no context!!!\n");
        return;
    }
    context->pt_regs->tp = (uint64_t)current_task;
    /* set minimal pt_regs fields (epc, sp) */
    context->pt_regs->epc = entry;

    context->pt_regs->sp = stack;

    /* a0 already set to 0 in init; clear other caller-saved regs for safety */
    context->pt_regs->a0 = 0;

    uint64_t old_sstatus;
    __asm__ volatile("csrr %0, sstatus" : "=r"(old_sstatus));  // 读取当前 sstatus
    
    context->pt_regs->sstatus = (old_sstatus & ~((1UL << 8) | (1UL << 1) | (1UL << 18))) | (1UL << 5);

}

void arch_to_user_mode(arch_context_t *context, uint64_t entry, uint64_t stack)
{   
    arch_context_to_user_mode(context, entry, stack);
    
    __asm__ __volatile__ (
        /* 使用临时寄存器清除 SUM（bit 18） */
        "li t0, %0\n\t"
        "csrc sstatus, t0\n\t"
        :
        : "i"(1UL << 18)
        : "t0", "memory"
    );
    printk("===================================================\n");
    uint64_t satp_val;
    __asm__ volatile("csrr %0, satp" : "=r"(satp_val));
    printk("SATP = 0x%lx\n", (unsigned long)satp_val);
    printk("===================================================\n");
    __asm__ __volatile__(
        "mv sp, %0\n\t"
        "j ret_from_trap_handler\n\t"
        :
        : "r"(context->pt_regs)
        : "memory");
}

void arch_yield()
{
    asm volatile("ecall" ::: "memory");
}

#define ARCH_SET_GS 0x1001
#define ARCH_SET_FS 0x1002
#define ARCH_GET_FS 0x1003
#define ARCH_GET_GS 0x1004

uint64_t sys_arch_prctl(uint64_t cmd, uint64_t arg)
{
    // RISC-V 没有 FS/GS 段寄存器，保持接口一致但不做实际操作
    (void)arg;

    switch (cmd)
    {
    case ARCH_SET_FS:
    case ARCH_SET_GS:
        // 忽略设置请求
        return 0;
    case ARCH_GET_FS:
    case ARCH_GET_GS:
        // 返回 0 表示未设置
        return 0;
    default:
        return (uint64_t)(-ENOSYS);
    }
}
